#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image.h"
#include "stb_image_write.h"
#include <iostream>
#include <cuda_runtime.h>

#define CHANNELS 3  // Number of color channels (RGB)

// CUDA kernel function for mean filter
__global__ void meanFilter(unsigned char* input, unsigned char* output, int width, int height) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= width || y >= height) return;

    int pixelIndex = (y * width + x) * CHANNELS;
    int filterSize = 10;
    int halfFilterSize = filterSize / 2;

    int r = 0, g = 0, b = 0;
    int count = 0;

    // Apply the mean filter
    for (int fy = -halfFilterSize; fy <= halfFilterSize; fy++) {
        for (int fx = -halfFilterSize; fx <= halfFilterSize; fx++) {
            // 相当于上课ppt中的
            /**
             if(curRow > -1 && curRow < h && curCol > -1 && curCol < w) {
                pixVal += in[curRow * w + curCol]
            }
             */
            int ix = min(max(x + fx, 0), width - 1);
            int iy = min(max(y + fy, 0), height - 1);
            int iPixelIndex = (iy * width + ix) * CHANNELS;

            r += input[iPixelIndex];
            g += input[iPixelIndex + 1];
            b += input[iPixelIndex + 2];
            count++;
        }
    }

    output[pixelIndex] = r / count;
    output[pixelIndex + 1] = g / count;
    output[pixelIndex + 2] = b / count;
}

void checkCudaError(cudaError_t err, const char* message) {
    if (err != cudaSuccess) {
        std::cerr << "Error: " << message << " (" << cudaGetErrorString(err) << ")" << std::endl;
        exit(EXIT_FAILURE);
    }
}

int main() {
    int width, height, channels;

    // Load input image using stb_image
    unsigned char* h_input = stbi_load("input.jpg", &width, &height, &channels, CHANNELS);
    if (!h_input) {
        std::cerr << "Error: Could not open or find the image." << std::endl;
        return -1;
    }

    // Allocate host memory for output
    unsigned char* h_output = new unsigned char[width * height * CHANNELS];

    // Allocate device memory
    unsigned char* d_input;
    unsigned char* d_output;
    checkCudaError(cudaMalloc(&d_input, width * height * CHANNELS), "cudaMalloc d_input");
    checkCudaError(cudaMalloc(&d_output, width * height * CHANNELS), "cudaMalloc d_output");

    // Copy input data from host to device
    checkCudaError(cudaMemcpy(d_input, h_input, width * height * CHANNELS, cudaMemcpyHostToDevice), "cudaMemcpy h_input to d_input");

    // Define grid and block dimensions
    // 2-dimension
    dim3 blockDim(16, 16);
    dim3 gridDim((width + blockDim.x - 1) / blockDim.x, (height + blockDim.y - 1) / blockDim.y);

    // Launch the CUDA kernel
    meanFilter<<<gridDim, blockDim>>>(d_input, d_output, width, height);

    // Check for any kernel launch errors
    checkCudaError(cudaGetLastError(), "Kernel launch failed");

    // Copy output data from device to host
    checkCudaError(cudaMemcpy(h_output, d_output, width * height * CHANNELS, cudaMemcpyDeviceToHost), "cudaMemcpy d_output to h_output");

    // Save output image using stb_image_write
    if (!stbi_write_jpg("output.jpg", width, height, CHANNELS, h_output, 100)) {
        std::cerr << "Error: Could not write the output image." << std::endl;
    }

    // Free device memory
    checkCudaError(cudaFree(d_input), "cudaFree d_input");
    checkCudaError(cudaFree(d_output), "cudaFree d_output");

    // Free host memory
    stbi_image_free(h_input);
    delete[] h_output;

    std::cout << "Image processing completed." << std::endl;

    return 0;
}
